import torch
import scipy.io as scio
import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
import matplotlib.pyplot as plt
from multiprocessing import Process
import numpy as np
from DSPPytorch import *
import util
import scipy.io as scio

BASE_DIR = os.path.dirname(__file__)


def genDataset(lp_range=[1]):
    lp_range = [*lp_range]
    pick_sym_num = 131072
    modOrder = 4
    datasetSaveDir = os.path.join(
        BASE_DIR, 'data/simulation/16QAM32Gbaud3200kmHe/PMD_eliminated_by_CMA'
    )
    if not os.path.exists(datasetSaveDir):
        os.makedirs(datasetSaveDir)
        print(f'Create dataset directory at \"{datasetSaveDir}\"')

    for lpIndx, lp in enumerate(lp_range):
        '''________________________________TraingSet__________________________________'''
        dataPath = os.path.join(
            BASE_DIR, 'data/simulation/16QAM32Gbaud3200kmHe/BPS_4_seed_1_dBm_{lp:d}_Loops_8.mat'.format(lp=lp)
        )
        '''读取文件信息'''
        data = scio.loadmat(dataPath)
        symbolRate = 32e9
        bitsPerSymbol = data['BitsPerSymbol'][0, 0]
        spanNum = data['Loops_N'][0, 0] * 5
        rolloff = 0.01
        lp = data['dBm_N'][0, 0]
        prbsx = data['prbsx']
        prbsy = data['prbsy']
        spanLen = 80e3
        L = spanLen * spanNum
        D = 17e-6
        DL_remain = D * L / 2  # 补偿了50%的色散
        sigx = torch.from_numpy(data['signalx'])
        sigy = torch.from_numpy(data['signaly'])
        sigx = sigx.reshape(1, 1, -1)
        sigy = sigy.reshape(1, 1, -1)
        sig = torch.cat([sigx, sigy], dim=1)

        sigxb2b = torch.from_numpy(data['signalx_label'])
        sigyb2b = torch.from_numpy(data['signaly_label'])
        sigxb2b = sigxb2b.reshape(1, 1, -1)
        sigyb2b = sigyb2b.reshape(1, 1, -1)
        sigb2b = torch.cat([sigxb2b, sigyb2b], dim=1)

        constellations = torch.from_numpy(util.CONST_16QAM)

        '''实例化后端的DSP类'''
        # mf = MatchedFilterLayer(alpha=rolloff, span=100, sample_factor=8, case_num=2)  # span: Correlated symbols
        pre_edc = EDCLayer(symbolRate=symbolRate, DL=DL_remain, case_num=2, tap=4096)
        post_edc = EDCLayer(symbolRate=symbolRate, DL=-DL_remain, case_num=2, tap=4096)
        pr = PhaseRecLayer(1024)
        lms = FIRLayer(tap=32, case_num=2, power_norm=True, centor_one=True)

        '''DSP procedure for the signal through the link'''
        # sig = mf(sig)
        sig = sig[..., 1::4]  # resampling
        sig = sig[..., 0:pick_sym_num * 2]
        prbsx = prbsx[..., 0:pick_sym_num * modOrder]
        prbsy = prbsy[..., 0:pick_sym_num * modOrder]
        prbs = np.concatenate([prbsx, prbsy], axis=0)

        if torch.cuda.is_available():
            sig = sig.cuda()
            lms = lms.cuda()
            pre_edc = pre_edc.cuda()
            post_edc = post_edc.cuda()

        with torch.no_grad():
            sig = pre_edc(sig)

            '''后端的TDE均衡PMD以及PDL等，是自适应的方法'''
            sig = sig.cpu()
            sig = FIRLayer.time_vary_infer(sig, err_mode='Godard', iter_num=2, lr=5e-4, tap=13)

            sig_np = pick_pol(sig.data.cpu().numpy().squeeze(), prbs, sample_factor=2)
            sig = torch.from_numpy(sig_np.copy())
            sig = sig[np.newaxis, ...]

            if torch.cuda.is_available():
                sig = sig.cuda()  # 重新放回cuda上

            _, _, good_rotate = util.pr_ber(sig[0, :, 1::2].cpu().data.numpy(), prbs, constellations.cpu().data.numpy())
            pr = PhaseRecLayer(1024, torch.from_numpy(good_rotate))
            '''检查一下偏振态对不对'''
            sig_check = sig[..., 1::2]
            sig_check = pr(sig_check)

            util.colorfulPlot(sig_check[0, 0, :].cpu().data.numpy(), prbs[0, ...], show=False,
                              saveFig=True, savePath=os.path.join(datasetSaveDir, 'constellations_check_result/trSet_constellation_lp_{lp:d}_x'.format(lp=lp)))
            util.colorfulPlot(sig_check[0, 1, :].cpu().data.numpy(), prbs[1, ...], show=False,
                              saveFig=True, savePath=os.path.join(datasetSaveDir, 'constellations_check_result/trSet_constellation_lp_{lp:d}_y'.format(lp=lp)))
            '''恢复原有的色散！'''
            sig = post_edc(sig)


        sig = sig.cpu().data.numpy().squeeze()
        labelx = util.map(prbs[0, ...], constellations)
        labely = util.map(prbs[1, ...], constellations)
        labelx = labelx.reshape(1, -1)
        labely = labely.reshape(1, -1)
        label = np.concatenate([labelx, labely], axis=0)

        save_dict = {
            'lp': lp,
            'symbolRate': symbolRate,
            'spanNum': spanNum,
            'spanLen': spanLen,
            'prbs': prbs,
            'sig': sig,
            'sig_label': label,
            'phase_state': good_rotate,
        }
        scio.savemat(os.path.join(datasetSaveDir, 'trSet_lp_{lp:d}.mat'.format(lp=lp)), save_dict)
        print(f'LP {lp} training set done!')
        del data
        '''_____________________Training set end____________________________________'''

        '''______________________Test set___________________________________________'''
        dataPath = os.path.join(
            BASE_DIR, 'data/simulation/16QAM32Gbaud3200kmHe/BPS_4_seed_2_dBm_{lp:d}_Loops_8.mat'.format(lp=lp)
        )
        '''读取文件信息'''
        data = scio.loadmat(dataPath)
        symbolRate = 32e9
        bitsPerSymbol = data['BitsPerSymbol'][0, 0]
        spanNum = data['Loops_N'][0, 0] * 5
        rolloff = 0.01
        lp = data['dBm_N'][0, 0]
        prbsx = data['prbsx']
        prbsy = data['prbsy']
        spanLen = 80e3
        L = spanLen * spanNum
        D = 17e-6
        DL_remain = D * L / 2  # 补偿了50%的色散
        sigx = torch.from_numpy(data['signalx'])
        sigy = torch.from_numpy(data['signaly'])
        sigx = sigx.reshape(1, 1, -1)
        sigy = sigy.reshape(1, 1, -1)
        sig = torch.cat([sigx, sigy], dim=1)

        sigxb2b = torch.from_numpy(data['signalx_label'])
        sigyb2b = torch.from_numpy(data['signaly_label'])
        sigxb2b = sigxb2b.reshape(1, 1, -1)
        sigyb2b = sigyb2b.reshape(1, 1, -1)
        sigb2b = torch.cat([sigxb2b, sigyb2b], dim=1)

        constellations = torch.from_numpy(util.CONST_16QAM)

        '''实例化后端的DSP类'''
        # mf = MatchedFilterLayer(alpha=rolloff, span=100, sample_factor=8, case_num=2)  # span: Correlated symbols
        pre_edc = EDCLayer(symbolRate=symbolRate, DL=DL_remain, case_num=2, tap=4096)
        post_edc = EDCLayer(symbolRate=symbolRate, DL=-DL_remain, case_num=2, tap=4096)
        pr = PhaseRecLayer(1024)
        lms = FIRLayer(tap=32, case_num=2, power_norm=True, centor_one=True)

        '''DSP procedure for the signal through the link'''
        # sig = mf(sig)
        sig = sig[..., 1::4]  # resampling
        sig = sig[..., 0:pick_sym_num * 2]
        prbsx = prbsx[..., 0:pick_sym_num * modOrder]
        prbsy = prbsy[..., 0:pick_sym_num * modOrder]
        prbs = np.concatenate([prbsx, prbsy], axis=0)

        if torch.cuda.is_available():
            sig = sig.cuda()
            lms = lms.cuda()
            pre_edc = pre_edc.cuda()
            post_edc = post_edc.cuda()

        with torch.no_grad():
            sig = pre_edc(sig)

            '''后端的TDE均衡PMD以及PDL等，是自适应的方法'''
            sig = sig.cpu()
            sig = FIRLayer.time_vary_infer(sig, err_mode='Godard', iter_num=2, lr=5e-4, tap=13)

            sig_np = pick_pol(sig.data.cpu().numpy().squeeze(), prbs, sample_factor=2)
            sig = torch.from_numpy(sig_np.copy())
            sig = sig[np.newaxis, ...]

            if torch.cuda.is_available():
                sig = sig.cuda()  # 重新放回cuda上

            _, _, good_rotate = util.pr_ber(sig[0, :, 1::2].cpu().data.numpy(), prbs, constellations.cpu().data.numpy())
            pr = PhaseRecLayer(1024, torch.from_numpy(good_rotate))
            '''检查一下偏振态对不对'''
            sig_check = sig[..., 1::2]
            sig_check = pr(sig_check)

            util.colorfulPlot(sig_check[0, 0, :].cpu().data.numpy(), prbs[0, ...], show=False,
                              saveFig=True, savePath=os.path.join(datasetSaveDir,
                                                                  'constellations_check_result/tstSet_constellation_lp_{lp:d}_x'.format(
                                                                      lp=lp)))
            util.colorfulPlot(sig_check[0, 1, :].cpu().data.numpy(), prbs[1, ...], show=False,
                              saveFig=True, savePath=os.path.join(datasetSaveDir,
                                                                  'constellations_check_result/tstSet_constellation_lp_{lp:d}_y'.format(
                                                                      lp=lp)))
            '''恢复原有的色散！'''
            sig = post_edc(sig)

        sig = sig.cpu().data.numpy().squeeze()
        labelx = util.map(prbs[0, ...], constellations)
        labely = util.map(prbs[1, ...], constellations)
        labelx = labelx.reshape(1, -1)
        labely = labely.reshape(1, -1)
        label = np.concatenate([labelx, labely], axis=0)

        save_dict = {
            'lp': lp,
            'symbolRate': symbolRate,
            'spanNum': spanNum,
            'spanLen': spanLen,
            'prbs': prbs,
            'sig': sig,
            'sig_label': label,
            'phase_state': good_rotate,
        }
        scio.savemat(os.path.join(datasetSaveDir, 'tstSet_lp_{lp:d}.mat'.format(lp=lp)), save_dict)
        print(f'LP {lp} test set done!')
        del data
        '''______________________Test set end____________________________________________'''

if __name__ == '__main__':

    '''这个程序是用来对数据进行预处理的，方便后面对训练集和测试集直接调用，在这个例子中，首先补偿了色散
    ，然后补偿了PMD，确定了相位旋转的大小, 之后又将色散补偿回来了！'''
    genDataset(lp_range=[-2, -1, 0, 1, 2, 3, 4])